import torch
import torchvision
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import os
import numpy as np

data_dir = './data'
model_dir = ['./checkpoint/B0',
             './checkpoint/B1',
             './checkpoint/B2',
             './checkpoint/B3',
             './checkpoint/FcaB1',
             './checkpoint/ShakeDropB1',
             './checkpoint/ablationstudy']

# 图像预处理
batch_size = 64  # 原文是1024
num_workers = 4  # 2
img_size = 128


def gen_mean_std(dataset):
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=2)
    train = next(iter(dataloader))[0]
    mean = np.mean(train.numpy(), axis=(0, 2, 3))
    std = np.std(train.numpy(), axis=(0, 2, 3))
    return mean, std


# if __name__ == "__main__":
#     train_dataset = torchvision.datasets.CIFAR100(root='data/',
#                                                   train=True,
#                                                   transform=transforms.Compose([transforms.ToTensor()]),
#                                                   download=False)
#     a, b = gen_mean_std(train_dataset)
#     print(a,b)


mean = [0.523129, 0.507429, 0.47282392]  # 128[0.49729094, 0.48025268, 0.42963162]
std = [0.26472256, 0.2599145, 0.28342927]  # 128[0.26493478, 0.2575264, 0.27690575]

# 原文提到了random erase的处理方法 但是有网页指出 这个方法对CIFAR-100没什么帮助
# 参考网页：https://blog.csdn.net/u013685264/article/details/122564323
transform_train = transforms.Compose([
    transforms.Resize((img_size, img_size)),
    # transforms.RandomCrop(64, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize(mean, std),
    transforms.RandomErasing()
])
transform_test = transforms.Compose(
    [transforms.Resize((128, 128)),
     transforms.ToTensor(),
     transforms.Normalize(mean, std)])

# CIFAR-100 数据集下载

data_name = './data/cifar-100-python'
if not os.path.isdir(data_name):
    train_dataset = torchvision.datasets.CIFAR100(root='data/',
                                                  train=True,
                                                  transform=transform_train,
                                                  download=True)
else:
    train_dataset = torchvision.datasets.CIFAR100(root='data/',
                                                  train=True,
                                                  transform=transform_train,
                                                  download=False)

test_dataset = torchvision.datasets.CIFAR100(root='data/',
                                             train=False,
                                             transform=transform_test)

# 数据载入

train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           num_workers=num_workers,
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          num_workers=num_workers,
                                          shuffle=False)
